package org.korosoft.rusalad.syncnrun.command;

import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.korosoft.rusalad.syncnrun.messaging.Message;
import org.korosoft.rusalad.syncnrun.messaging.MessageProcessor;
import org.korosoft.rusalad.syncnrun.messaging.MessageRoute;
import org.korosoft.rusalad.syncnrun.messaging.StateAwareMessageProcessor;
import org.korosoft.rusalad.syncnrun.tools.BlockingQueueExecutor;
import org.korosoft.rusalad.syncnrun.tools.Utils;

import java.io.*;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.security.MessageDigest;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Syncnrun "upload files" and "download files" RPC commands common ancestor.
 *
 * @author Dmitry Korotkov
 * @since 1.0
 */
public abstract class AbstractSyncCmd<T extends Answer> implements MessageDrivenCommand<T> {
    private static final Log log = LogFactory.getLog(AbstractSyncCmd.class);
    public static final int CHUNK_SIZE = 65536;
    public static final String ALGORITHM = "SHA-256";
    public static final int MAX_HASHES_PER_LIST = 512; // 64K * 512 = 32M of raw data = 16K of hashes

    protected AbstractSyncCmd() {
        initTransients();
    }

    /*
    ================================= Public API =================================
     */

    public static class FileInfo implements Serializable {
        private final String name;
        private final long size;
        private final long lastModified;

        public FileInfo(String name, long size, long lastModified) {
            this.name = name;
            this.size = size;
            this.lastModified = lastModified;
        }

        public String toString() {
            return name + " (" + size + " bytes)";
        }
    }

    public static class LocalFilesInfo {
        public final List<AbstractSyncCmd.FileInfo> localFiles;
        public final File localFilesPath;
        public final boolean singleFileTransfer;

        public LocalFilesInfo(List<FileInfo> localFiles, File localFilesPath, boolean singleFileTransfer) {
            this.localFiles = localFiles;
            this.localFilesPath = localFilesPath;
            this.singleFileTransfer = singleFileTransfer;
        }
    }

    public static AbstractSyncCmd.LocalFilesInfo prepareLocalFilesInfo(String source) {
        List<AbstractSyncCmd.FileInfo> localFiles = new ArrayList<AbstractSyncCmd.FileInfo>();

        log.debug("Scanning local files...");

        File file = new File(source);
        if (!file.exists()) {
            throw new RuntimeException("Source file " + file + " does not exist");
        }
        int origPathLength = file.getAbsolutePath().length();

        final File sourcePath;
        if (file.isDirectory()) {
            walkFiles(localFiles, file, origPathLength);
            sourcePath = file;
        } else {
            localFiles.add(new AbstractSyncCmd.FileInfo(file.getName(), FileUtils.sizeOf(file), file.lastModified()));
            sourcePath = file.getParentFile();
        }

        if (log.isDebugEnabled()) {
            log.debug(String.format("Found %s", getFileStatsString(localFiles)));
        }

        return new AbstractSyncCmd.LocalFilesInfo(localFiles, sourcePath, !file.isDirectory());
    }

    /*
    ================================= Ancestor API =================================
     */

    protected void processHashListMessage(MessageRoute route, FileHashListMsg message, final List<AbstractSyncCmd.FileInfo> sourceFiles, final File sourcePath, final Runnable onFinishCallback) {
        try {// ToDo: update
            byte[] chunk = new byte[CHUNK_SIZE];

            // Process all files for which we received chunk lists first
            for (FileHashListMsg.FileHash hash : message.allFiles) {
                File file = new File(sourcePath, sourceFiles.get(hash.fileId).name);
                processedFlags.put(hash.fileId, hash.fileId);
                chunksSent.addAndGet(processFile(route, chunk, bytesSent, hash, file));
            }
            if (message.numberOfMessagesSent > 0) {
                chunkListMessagesProcessed.addAndGet(message.numberOfMessagesSent);
            }
            if (chunkListMessagesProcessed.decrementAndGet() == 0) {
                // Process all files to be sent for which we didn't received chunk list (send complete files)
                for (int i = 0; i < sourceFiles.size(); i++) {
                    if (processedFlags.containsKey(i)) {
                        continue;
                    }
                    File file = new File(sourcePath, sourceFiles.get(i).name);
                    FileHashListMsg.FileHash hash = new FileHashListMsg.FileHash(i, Collections.<byte[]>emptyList());
                    chunksSent.addAndGet(processFile(route, chunk, bytesSent, hash, file));
                }
                route.put(new AllChunksAreSentMsg(chunksSent.get(), bytesSent.get()));
                if (onFinishCallback != null) {
                    onFinishCallback.run();
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    protected long sendOutHashes(MessageRoute route, final List<AbstractSyncCmd.FileInfo> sourceFiles, final String targetPath, final boolean singleFileTransfer) throws IOException, InterruptedException {
        final CountDownLatch doneLatch = new CountDownLatch(1);
        final AtomicInteger chunksReceived = new AtomicInteger(0);
        final AtomicLong bytesSent = new AtomicLong();

        if (log.isDebugEnabled()) {
            log.debug(String.format("Preparing directory structure for %s", getFileStatsString(sourceFiles)));
        }

        long totalBytes = 0;
        for (FileInfo fileInfo : sourceFiles) {
            totalBytes += fileInfo.size;
        }
        this.totalBytes.set(totalBytes);

        prepareTargetDirectory(sourceFiles, targetPath, singleFileTransfer);
        route.setHotReceiver(createServerReceiver(bytesSent, doneLatch, chunksReceived, sourceFiles, targetPath));
        hashExistingFiles(sourceFiles, targetPath, route);
        doneLatch.await();
        for (FileInfo fileInfo : sourceFiles) {
            File f = new File(targetPath, fileInfo.name);
            if (!f.setLastModified(fileInfo.lastModified)) {
                log.warn("Failed to set last modified time for file " + f.getAbsolutePath());
            }
        }
        return bytesSent.get();
    }

    protected static class FileHashListMsg implements Message {
        private final List<FileHash> allFiles;
        /**
         * May be -1 when number is undefined
         */
        private final int numberOfMessagesSent;

        private FileHashListMsg(List<FileHash> allFiles, int numberOfMessagesSent) {
            this.allFiles = allFiles;
            this.numberOfMessagesSent = numberOfMessagesSent;
        }

        private static class FileHash implements Serializable {
            private final int fileId;
            private final List<byte[]> hashes;

            private FileHash(int fileId, List<byte[]> hashes) {
                this.fileId = fileId;
                this.hashes = hashes;
            }

            @Override
            public String toString() {
                return "File Id: " + fileId + " (" + hashes.size() + " hashes)";
            }
        }

        @Override
        public String toString() {
            return String.format("File hash list. File hashes: %s; Number of messages sent: %s", Arrays.toString(allFiles.toArray()), numberOfMessagesSent == -1 ? "unavailable" : Integer.toString(numberOfMessagesSent));
        }
    }

    /*
    ================================= Private implementation =================================
     */

    private transient AtomicInteger chunkListMessagesProcessed;
    private transient ConcurrentMap<Integer, Object> processedFlags;
    private transient AtomicInteger chunksSent;
    private transient AtomicLong bytesSent;
    private transient AtomicLong bytesHashed;
    private transient AtomicLong bytesTransferred;
    private transient AtomicLong totalBytes;
    private transient LinkedBlockingQueue<TransferStatsRecord> transferStatsRecords;
    private transient AtomicLong lastStatsUpdate;

    protected void initTransients() {
        chunkListMessagesProcessed = new AtomicInteger(0);
        processedFlags = new ConcurrentHashMap<Integer, Object>();
        chunksSent = new AtomicInteger(0);
        bytesSent = new AtomicLong(0);
        bytesHashed = new AtomicLong(0);
        bytesTransferred = new AtomicLong(0);
        totalBytes = new AtomicLong(0);
        transferStatsRecords = new LinkedBlockingQueue<TransferStatsRecord>();
        lastStatsUpdate = new AtomicLong(0);
    }


    private void readObject(ObjectInputStream ois)
            throws ClassNotFoundException, IOException {
        ois.defaultReadObject();
        // Read/initialize additional fields
        initTransients();
    }

    private static void walkFiles(List<AbstractSyncCmd.FileInfo> localFiles, File file, int origPathLength) {
        File[] filesList = file.listFiles();
        if (filesList == null) {
            return;
        }
        for (File f : filesList) {
            if (f.isDirectory()) {
                localFiles.add(new AbstractSyncCmd.FileInfo(f.getAbsolutePath().substring(origPathLength) + "/", 0, f.lastModified()));
                walkFiles(localFiles, f, origPathLength);
            } else {
                localFiles.add(new AbstractSyncCmd.FileInfo(f.getAbsolutePath().substring(origPathLength) + "", FileUtils.sizeOf(f), f.lastModified()));
            }
        }
    }

    private MessageProcessor createServerReceiver(final AtomicLong bytesSent, final CountDownLatch doneLatch, final AtomicInteger chunksReceived, final List<AbstractSyncCmd.FileInfo> sourceFiles, final String targetPath) {
        final ExecutorService singleThreadDiskWriter = Executors.newSingleThreadExecutor(new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                return new Thread(r, "disk-writer");
            }
        });
        final Map<File, RandomAccessFile> filesCache = new HashMap<File, RandomAccessFile>();

        return new StateAwareMessageProcessor() {
            @Override
            public void transferEnded(MessageRoute route) {
                final CountDownLatch latch = new CountDownLatch(1);
                singleThreadDiskWriter.execute(new Runnable() {
                    @Override
                    public void run() {
                        for (RandomAccessFile randomAccessFile : filesCache.values()) {
                            IOUtils.closeQuietly(randomAccessFile);
                        }
                        filesCache.clear();
                        latch.countDown();
                    }
                });
                try {
                    latch.await();
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }

            @Override
            public void processMessage(final MessageRoute route, final Message message) {
                log.trace("Received " + message.toString());
                singleThreadDiskWriter.execute(new Runnable() {
                    @Override
                    public void run() {
                        if (message instanceof AbstractSyncCmd.AllChunksAreSentMsg) {
                            bytesSent.set(((AllChunksAreSentMsg) message).bytesSent);
                            if (chunksReceived.addAndGet(((AbstractSyncCmd.AllChunksAreSentMsg) message).chunksSent) == 0) {
                                done();
                            }
                            return;
                        }
                        if (message instanceof AbstractSyncCmd.FileChunkMsg) {
                            AbstractSyncCmd.FileChunkMsg m = (AbstractSyncCmd.FileChunkMsg) message;

                            updateTransferStats(m);

                            File file = new File(targetPath, sourceFiles.get(m.fileId).name);

                            try {
                                RandomAccessFile randomAccessFile = filesCache.get(file);
                                if (randomAccessFile == null) {
                                    randomAccessFile = new RandomAccessFile(file, "rw");
                                    filesCache.put(file, randomAccessFile);
                                }
                                FileChannel channel = randomAccessFile.getChannel();
                                channel.write(ByteBuffer.wrap(m.data), m.offset);
                            } catch (IOException e) {
                                throw new RuntimeException(e);
                            }

                            if (chunksReceived.decrementAndGet() == 0) {
                                done();
                            }

                            return;
                        }
                        throw new RuntimeException(String.format("Unexpected message: %s", message.getClass().getName()));
                    }
                });
            }

            private void done() {
                for (RandomAccessFile randomAccessFile : filesCache.values()) {
                    IOUtils.closeQuietly(randomAccessFile);
                }
                filesCache.clear();
                doneLatch.countDown();
            }
        };
    }

    private void updateTransferStats(FileChunkMsg m) {
        if (m != null) {
            bytesTransferred.addAndGet(m.data.length);
            transferStatsRecords.add(new TransferStatsRecord(m.data.length));
        }
        final long t = System.currentTimeMillis();
        for (; ; ) {
            TransferStatsRecord r = transferStatsRecords.peek();
            if (r == null || (t - r.timeMillis) < 60000) {
                break;
            }
            transferStatsRecords.poll();
        }
        if (System.currentTimeMillis() - lastStatsUpdate.get() > 1000) {
            lastStatsUpdate.set(t);
            long transferRate = 0;
            for (TransferStatsRecord r : transferStatsRecords) {
                transferRate += r.bytes;
            }
            TransferStatsRecord r = transferStatsRecords.peek();
            if (r != null) {
                long secondsElapsed = (t - r.timeMillis) / 1000;
                if (secondsElapsed != 0) {
                    transferRate = transferRate / secondsElapsed; // Get transfer rate per second
                } else {
                    transferRate = 0;
                }
            }

            if (bytesHashed.get() == 0) {
                log.debug(String.format("Transfer stats: %s transferred, transfer rate: %s/s",
                        Utils.getHumanReadableDataVolume(bytesTransferred.get()),
                        Utils.getHumanReadableDataVolume(transferRate)));
            } else {
                log.debug(String.format("Transfer stats: %s out of %s hashed, %s transferred, transfer rate: %s/s",
                        Utils.getHumanReadableDataVolume(bytesHashed.get()), Utils.getHumanReadableDataVolume(totalBytes.get()),
                        Utils.getHumanReadableDataVolume(bytesTransferred.get()),
                        Utils.getHumanReadableDataVolume(transferRate)));
            }
        }
    }

    private int processFile(final MessageRoute route, byte[] chunk, final AtomicLong bytesSent, final FileHashListMsg.FileHash hash, File file) throws IOException, InterruptedException {
        if (file.isDirectory()) {
            return 0;
        }
        final AtomicInteger chunksSent = new AtomicInteger(0);
        final InputStream stream = new BufferedInputStream(new FileInputStream(file));
        int chunkNo = 0;
        long offset = 0;
        final AtomicInteger chunksInProgress = new AtomicInteger(1);
        final CountDownLatch allChunkAreHashed = new CountDownLatch(1);
        for (int l = readChunk(chunk, stream); l > 0; l = readChunk(chunk, stream)) {
            final byte[] chunkCopy = Arrays.copyOf(chunk, l);
            final int chunkNoFinal = chunkNo;
            final long offsetFinal = offset;
            chunksInProgress.incrementAndGet();
            hashTaskExecutor.addHashTask(new HashTask(chunkCopy) {
                @Override
                protected void done(byte[] h) throws InterruptedException {
                    if (chunkNoFinal >= hash.hashes.size() || !Arrays.equals(h, hash.hashes.get(chunkNoFinal))) {
                        // Send chunk
                        FileChunkMsg message = new FileChunkMsg(hash.fileId, chunkCopy, offsetFinal);
                        route.put(message);
                        chunksSent.incrementAndGet();
                        bytesSent.addAndGet(chunkCopy.length);
                        updateTransferStats(message);
                    }
                    if (chunksInProgress.decrementAndGet() == 0) {
                        allChunkAreHashed.countDown();
                    }
                }
            });
            chunkNo++;
            offset += l;
        }
        if (chunksInProgress.decrementAndGet() == 0) {
            allChunkAreHashed.countDown();
        }
        allChunkAreHashed.await();
        stream.close();
        return chunksSent.get();
    }

    private abstract static class HashTask {
        private final byte[] chunk;

        private HashTask(byte[] chunk) {
            this.chunk = chunk;
        }

        protected abstract void done(byte[] hash) throws InterruptedException;
    }

    private static class HashTaskExecutor {
        private final ThreadLocal<MessageDigest> messageDigestStore = new ThreadLocal<MessageDigest>();
        private final ExecutorService executorService = new BlockingQueueExecutor("hash-calculator-", 64, Math.max(1, Runtime.getRuntime().availableProcessors()));

        public void addHashTask(final HashTask hashTask) {
            executorService.execute(new Runnable() {
                @Override
                public void run() {
                    try {
                        MessageDigest messageDigest = messageDigestStore.get();
                        if (messageDigest == null) {
                            messageDigest = MessageDigest.getInstance(ALGORITHM);
                            messageDigestStore.set(messageDigest);
                        }
                        messageDigest.update(hashTask.chunk);
                        byte[] digest = messageDigest.digest();
                        hashTask.done(Arrays.copyOf(digest, digest.length));
                    } catch (Exception e) {
                        log.error(e.getMessage());
                        log.trace(e.getMessage(), e);
                    }
                }
            });
        }
    }

    private static final HashTaskExecutor hashTaskExecutor = new HashTaskExecutor();

    private void hashExistingFiles(List<AbstractSyncCmd.FileInfo> sourceFiles, String targetPath, MessageRoute route) throws IOException, InterruptedException {
        // Hash existing files
        File targetDir = new File(targetPath);
        final byte[] chunk = new byte[CHUNK_SIZE];

        final List<FileHashListMsg.FileHash> allFiles = new ArrayList<FileHashListMsg.FileHash>(sourceFiles.size());
        int numberOfHashes = 0;
        int numberOfMessages = 0;
        for (int fileId = 0; fileId < sourceFiles.size(); fileId++) {
            final String name = sourceFiles.get(fileId).name;
            File file = new File(targetDir, name);
            if (!file.exists()) {
                log.info(String.format("File %s is missing. Requesting full transfer by not sending its hashes.", file));
                RandomAccessFile f = new RandomAccessFile(file, "rw");
                f.setLength(sourceFiles.get(fileId).size);
                f.close();
                bytesHashed.addAndGet(sourceFiles.get(fileId).size);
                continue;
            }
            final List<byte[]> hashes = Collections.synchronizedList(new ArrayList<byte[]>(FileUtils.sizeOf(file) / CHUNK_SIZE + ((FileUtils.sizeOf(file) % CHUNK_SIZE)) > 0 ? 1 : 0));
            if (file.isDirectory()) {
                bytesHashed.addAndGet(sourceFiles.get(fileId).size);
                continue;
            }
            log.info("File " + file + " exists. Hashing file...");
            final AtomicInteger chunksLeftToHash = new AtomicInteger(1);
            final CountDownLatch allChunksAreHashed = new CountDownLatch(1);
            final InputStream stream = new BufferedInputStream(new FileInputStream(file));
            try {
                for (int l = readChunk(chunk, stream); l > 0; l = readChunk(chunk, stream)) {
                    final int hashNoFinal = hashes.size();
                    hashes.add(null);
                    chunksLeftToHash.incrementAndGet();
                    hashTaskExecutor.addHashTask(new HashTask(Arrays.copyOf(chunk, l)) {
                        @Override
                        protected void done(byte[] hash) throws InterruptedException {
                            hashes.set(hashNoFinal, hash);
                            bytesHashed.addAndGet(chunk.length);
                            updateTransferStats(null);
                            if (chunksLeftToHash.decrementAndGet() == 0) {
                                allChunksAreHashed.countDown();
                            }
                        }
                    });

                    if ((long) (hashes.size()) * (long) CHUNK_SIZE > sourceFiles.get(fileId).size) {
                        break; // We do not need to hash more than remote file length.
                    }
                }
            } finally {
                stream.close();
            }
            if (chunksLeftToHash.decrementAndGet() == 0) {
                allChunksAreHashed.countDown();
            } else {
                allChunksAreHashed.await();
            }
            allFiles.add(new FileHashListMsg.FileHash(fileId, new ArrayList<byte[]>(hashes)));
            if (!file.canWrite()) {
                if (!file.setWritable(true)) {
                    log.warn(String.format("Failed to set file '%s' writable", file.getAbsolutePath()));
                }
            }
            RandomAccessFile f = new RandomAccessFile(file, "rw");
            f.setLength(sourceFiles.get(fileId).size);
            f.close();
            numberOfHashes += hashes.size();
            if (numberOfHashes > MAX_HASHES_PER_LIST) {
                route.put(new FileHashListMsg(new ArrayList<FileHashListMsg.FileHash>(allFiles), -1));
                allFiles.clear();
                numberOfMessages++;
                numberOfHashes = 0;
            }
        }
        route.put(new FileHashListMsg(allFiles, numberOfMessages + 1));
    }

    private int readChunk(byte[] chunk, InputStream stream) throws IOException {
        int l = 0;
        while (l < CHUNK_SIZE) {
            int r = stream.read(chunk, l, CHUNK_SIZE - l);
            if (r == -1) {
                break;
            }
            l += r;
        }
        return l;
    }

    private void prepareTargetDirectory(List<AbstractSyncCmd.FileInfo> sourceFiles, String targetPath, boolean singleFileTransfer) {
        File targetDir = new File(targetPath);
        if (targetDir.exists()) {
            Set<File> files = new HashSet<File>();
            for (AbstractSyncCmd.FileInfo fileInfo : sourceFiles) {
                File file = new File(targetDir, fileInfo.name);
                files.add(file);
            }
            if (targetDir.isDirectory()) {
                if (!singleFileTransfer) {
                    removeMissingFiles(targetDir, files);
                }
            } else {
                if (!targetDir.delete() && !targetDir.delete()) {
                    throw new RuntimeException("Destination folder is actually a file and it cannot be deleted");
                }
                if (!targetDir.mkdirs()) {
                    throw new RuntimeException("Failed to create target directory");
                }
            }
        } else {
            if (!targetDir.mkdirs()) {
                throw new RuntimeException("Failed to create target directory");
            }
        }
        // Create directories
        for (AbstractSyncCmd.FileInfo fileInfo : sourceFiles) {
            File file = new File(targetDir, fileInfo.name);
            if (fileInfo.name.endsWith("/")) {
                if (file.exists()) {
                    if (!file.isDirectory()) {
                        if (!file.delete() && !file.delete()) {
                            throw new RuntimeException("Failed to delete file " + file);
                        }
                        if (!file.mkdirs()) {
                            throw new RuntimeException("Failed to create directory " + file);
                        }
                        log.info("Created directory " + file);
                    }
                } else {
                    if (!file.mkdirs()) {
                        throw new RuntimeException("Failed to create directory " + file);
                    }
                    log.info("Created directory " + file);
                }
            }
        }
    }

    private static String getFileStatsString(List<FileInfo> localFiles) {
        long totalBytes = 0;
        int totalFiles = 0;
        for (FileInfo fileInfo : localFiles) {
            totalBytes += fileInfo.size;
            totalFiles++;
        }
        return String.format("%d files of %s total volume", totalFiles, Utils.getHumanReadableDataVolume(totalBytes));
    }

    private void removeMissingFiles(File targetDir, Set<File> files) {
        File[] filesList = targetDir.listFiles();
        if (filesList == null) {
            return;
        }
        for (File f : filesList) {
            if (f.isDirectory()) {
                removeMissingFiles(f, files);
            }
            if (!files.contains(f)) {
                if (!f.delete() && !f.delete()) {
                    if (f.isDirectory()) {
                        throw new RuntimeException("Failed to delete a directory " + f);
                    } else {
                        throw new RuntimeException("Failed to delete a file " + f);
                    }
                }
                log.info(String.format("Removed %s", f));
            }
        }
    }

    private static class AllChunksAreSentMsg implements Message {
        private final int chunksSent;
        private final long bytesSent;

        private AllChunksAreSentMsg(int chunksSent, long bytesSent) {
            this.chunksSent = chunksSent;
            this.bytesSent = bytesSent;
        }

        @Override
        public String toString() {
            return String.format("End of transfer message. Chunks sent: %d; Bytes sent: %d", chunksSent, bytesSent);
        }
    }

    private static class FileChunkMsg implements Message {
        private final int fileId;
        private final byte[] data;
        private final long offset;

        private FileChunkMsg(int fileId, byte[] data, long offset) {
            this.fileId = fileId;
            this.data = data;
            this.offset = offset;
        }

        @Override
        public String toString() {
            return String.format("Chunk of file %d at offset %d of length %d", fileId, offset, data.length);
        }
    }

    private static class TransferStatsRecord {
        private final long timeMillis;
        private final int bytes;

        private TransferStatsRecord(int bytes) {
            timeMillis = System.currentTimeMillis();
            this.bytes = bytes;
        }
    }

}
